import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os

# 设置更好的图表风格
plt.style.use('ggplot')

# 创建保存目录
save_dir = './dataset/channel_data/full_timeline'
os.makedirs(save_dir, exist_ok=True)

# 加载原始CSV数据
csv_path = './merged_output.csv'
print(f"Loading CSV data: {csv_path}")
df = pd.read_csv(csv_path)

# 获取所有时间点
time_points = sorted(df['receiver_point'].unique())
n_time_points = len(time_points)
print(f"Total time points: {n_time_points}")

# 特征名称和单位
feature_names = [
    'Horizontal Arrival Angle (degrees)', 
    'Vertical Arrival Angle (degrees)', 
    'Received Power (dB)', 
    'Horizontal Spread Angle (ASA)', 
    'Vertical Spread Angle (ZSA)'
]

# 特征列映射
feature_columns = [
    'horizontal_aoa',
    'vertical_aoa',
    'received_power',
    'ASA',
    'ZSA'
]

# 为每个时间点提取第一个簇(path_number=1)的数据
first_cluster_data = np.zeros((n_time_points, 5))  # [time_points, features]

print("Extracting first cluster data (path_number=1) for all time points...")
for t_idx, t in enumerate(time_points):
    # 获取当前时间点的数据
    time_data = df[df['receiver_point'] == t]
    
    # 按路径编号排序，选择第一个簇(path_number=1)
    sorted_paths = time_data.sort_values('path_number')
    first_cluster = sorted_paths.iloc[0]  # 第一个簇
    
    # 提取特征
    first_cluster_data[t_idx, 0] = first_cluster['horizontal_aoa']  # 水平到达角
    first_cluster_data[t_idx, 1] = first_cluster['vertical_aoa']    # 垂直到达角
    first_cluster_data[t_idx, 2] = first_cluster['received_power']  # 功率
    first_cluster_data[t_idx, 3] = first_cluster['ASA']             # 水平扩展角
    first_cluster_data[t_idx, 4] = first_cluster['ZSA']             # 垂直扩展角

# 为每个特征创建单独的图表
for feature_idx in range(5):
    feature_data = first_cluster_data[:, feature_idx]
    
    # 创建图表
    plt.figure(figsize=(15, 6))
    
    # 绘制折线图
    plt.plot(feature_data, '-', linewidth=1, alpha=0.7)
    
    # 根据需要添加散点图在其上
    # 如果点太多，只绘制部分点
    if n_time_points > 100:
        stride = n_time_points // 100
        plt.plot(range(0, n_time_points, stride), feature_data[::stride], 'o', markersize=4)
    else:
        plt.plot(feature_data, 'o', markersize=4)
    
    plt.title(f'First Cluster (Path #1): {feature_names[feature_idx]} Across All Time Points', fontsize=14)
    plt.xlabel('Time Point Index (1-{})'.format(n_time_points), fontsize=12)
    plt.ylabel(feature_names[feature_idx], fontsize=12)
    plt.grid(True)
    
    # 保存图表
    save_path = os.path.join(save_dir, f'first_cluster_{feature_columns[feature_idx]}_full.png')
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    
    print(f"Saved {feature_names[feature_idx]} plot to {save_path}")

# 创建功率变化的局部放大图，分段显示
segment_size = 200  # 每段显示的时间点数量
for start_idx in range(0, n_time_points, segment_size):
    end_idx = min(start_idx + segment_size, n_time_points)
    
    plt.figure(figsize=(15, 6))
    plt.plot(range(start_idx, end_idx), first_cluster_data[start_idx:end_idx, 2], 'o-', linewidth=1.5, markersize=4)
    plt.title(f'First Cluster (Path #1): Received Power (Time Points {start_idx+1}-{end_idx})', fontsize=14)
    plt.xlabel('Time Point Index', fontsize=12)
    plt.ylabel('Received Power (dB)', fontsize=12)
    plt.grid(True)
    
    save_path = os.path.join(save_dir, f'first_cluster_power_segment_{start_idx+1}_{end_idx}.png')
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    
    print(f"Saved power segment {start_idx+1}-{end_idx} plot")

# 创建归一化的特征比较图
plt.figure(figsize=(15, 8))

# 归一化所有特征
normalized_data = np.zeros_like(first_cluster_data)
for feature_idx in range(5):
    feature_data = first_cluster_data[:, feature_idx]
    min_val = feature_data.min()
    max_val = feature_data.max()
    if max_val > min_val:
        normalized_data[:, feature_idx] = (feature_data - min_val) / (max_val - min_val)
    else:
        normalized_data[:, feature_idx] = feature_data * 0

# 绘制所有归一化特征
for feature_idx in range(5):
    plt.plot(normalized_data[:, feature_idx], linewidth=1.5, label=feature_names[feature_idx].split("(")[0].strip())

plt.title('Normalized Metrics for First Cluster (Path #1) Across All Time Points', fontsize=16)
plt.xlabel('Time Point Index (1-{})'.format(n_time_points), fontsize=12)
plt.ylabel('Normalized Value', fontsize=12)
plt.legend(fontsize=10)
plt.grid(True)

# 保存归一化比较图
norm_save_path = os.path.join(save_dir, 'first_cluster_normalized_full.png')
plt.tight_layout()
plt.savefig(norm_save_path, dpi=150)
plt.close()

print(f"Saved normalized comparison to {norm_save_path}")

# 创建特征分布直方图
plt.figure(figsize=(15, 10))

for feature_idx in range(5):
    plt.subplot(2, 3, feature_idx + 1)
    feature_data = first_cluster_data[:, feature_idx]
    
    plt.hist(feature_data, bins=30, alpha=0.7)
    plt.title(f'Distribution of {feature_names[feature_idx].split("(")[0]}', fontsize=12)
    plt.xlabel(feature_names[feature_idx], fontsize=10)
    plt.ylabel('Frequency', fontsize=10)
    plt.grid(True)
    
    # 添加统计信息
    plt.text(0.05, 0.95, 
             f'Min: {feature_data.min():.2f}\nMax: {feature_data.max():.2f}\nMean: {feature_data.mean():.2f}\nStd: {feature_data.std():.2f}',
             transform=plt.gca().transAxes,
             verticalalignment='top',
             bbox=dict(boxstyle='round', facecolor='white', alpha=0.8),
             fontsize=9)

plt.suptitle('Distribution of First Cluster (Path #1) Metrics Across All Time Points', fontsize=16)
plt.tight_layout(rect=[0, 0, 1, 0.97])

# 保存分布图
dist_save_path = os.path.join(save_dir, 'first_cluster_distribution_full.png')
plt.savefig(dist_save_path, dpi=150)
plt.close()

print(f"Saved distribution plots to {dist_save_path}")
print("Full time span analysis complete!") 